#!/usr/bin/python

from google.appengine.ext import webapp
from cStringIO import StringIO
import logging
import simplejson as json
from protorpc.webapp import service_handlers


DEFAULT_REGISTRY_PATH = service_handlers.DEFAULT_REGISTRY_PATH


class ServiceException(Exception):
  """Base class for exceptions in apiserving."""


class BadRequestException(ServiceException):
  """Bad request exception that is mapped to a 400 response."""
  http_status = '400 BAD REQUEST'


class ForbiddenException(ServiceException):
  """Forbidden exception that is mapped to a 403 response."""
  http_status = '403 FORBIDDEN'


class InternalServerErrorException(ServiceException):
  """Internal server exception that is mapped to a 500 response."""
  http_status = '500 INTERNAL SEVER ERROR'


class NotFoundException(ServiceException):
  """Not found exception that is mapped to a 404 response."""
  http_status = '404 NOT FOUND'


class UnauthorizedException(ServiceException):
  """Unauthorized exception that is mapped to a 401 response."""
  http_status = '401 UNAUTHORIZED'


ERROR_NAME_MAP = {
  'bad request': BadRequestException,
  'forbidden': ForbiddenException,
  'internal server error': InternalServerErrorException,
  'not found': NotFoundException,
  'unauthorized': UnauthorizedException,
}


class ApiServer(object):
  """Simple WSGI application that wraps ProtoRPC, then reformats error messages.

  Usage:
  1) Create a apiserving.ApiServer instead of a webapp.WSGIApplication
  2) raise remote.ApplicationError("Custom message.", error_name="unauthorized")

  The allowed values for error_name are specified in ERROR_NAME_MAP above.
  The ServiceException classes specify the http status code returned.
  Below is a bit more detail on how to use it, followed by HTTP responses.

  - - - - services.py - - - -
  from google.appengine.ext.webapp import util
  import apiserving
  import postservice

  def main():
    util.run_wsgi_app(apiserving.ApiServer(
        [('/_ah/spi/Greetings_v1', postservice.PostService)], debug=True))

  if __name__ == '__main__':
    main()

  - - - - postservice.py - - - -
  class PostService(remote.Service):
    ...
    @remote.method(GetNotesRequest, Notes)
    def list(self, request):
      raise remote.ApplicationError("Please log in as an admin user.",
                                    error_name="unauthorized")

  - - - - ProtoRPC error format - - - -
  HTTP/1.0 400 Please log in as an admin user.
  content-type: application/json

  {
    "state": "APPLICATION_ERROR",
    "error_message": "Please log in as an admin user."
    "error_name": "unauthorized",
  }

  - - - - Reformatted error format - - - -
  HTTP/1.0 401 UNAUTHORIZED
  content-type: application/json

  {
    "error_message": "Please log in as an admin user."
  }

  """

  def __init__(self, services, registry_path=DEFAULT_REGISTRY_PATH, *args,
               **kwargs):
    mapping = service_handlers.service_mapping(services, registry_path)
    self.service_app = webapp.WSGIApplication(mapping, *args, **kwargs)

  def __call__(self, environ, start_response):

    def StartResponse(status, headers, exc_info=None):
      """Save args, defer start_response until response body is parsed.

      Create output buffer for body to be written into.
      Note: this is not quite WSGI compliant: The body should come back as an
        iterator returned from calling service_app() but instead, StartResponse
        returns a writer that will be later called to output the body.
      See google/appengine/ext/webapp/__init__.py::Response.wsgi_write()
          write = start_response('%d %s' % self.__status, self.__wsgi_headers)
          write(body)

      Args:
        status: Http status to be sent with this response
        headers: Http headers to be sent with this response
        exc_info: Exception info to be displayed for this response
      Returns:
        callable that takes as an argument the body content
      """
      call_context.update({'status': status, 'headers': headers,
                           'exc_info': exc_info, 'body_buffer': StringIO()})
      return call_context['body_buffer'].write

    call_context = {}
    body_iter = self.service_app(environ, StartResponse)
    if 'body_buffer' in call_context:
      body = call_context['body_buffer'].getvalue()
    if not body: # In case standard WSGI behavior is implemented later
      body = ''.join(body_iter)
    status, headers = call_context.get('status'), call_context.get('headers')
    if status[:3] == '400' and ('content-type', 'application/json') in headers:
      try:
        body_obj = json.loads(body)
      except ValueError:
        body_obj = {}
      if body_obj.get('state') == 'APPLICATION_ERROR':
        error_class = ERROR_NAME_MAP.get(body_obj.get('error_name'))
        if error_class:
          status = error_class.http_status
        replaced_body_obj = {'error_message': body_obj.get('error_message')}
        body = json.dumps(replaced_body_obj)
    start_response(status, headers, call_context.get('exc_info'))
    return [body]
